# eval_question2sql.py
# usage: 
# python ./src/tools/eval_question2sql.py --questions ./data/samples.round01.jsonl

import argparse
import json
from typing import List, Dict, Set
from clean_question import clean_question


def load_tmp_questions(file_path: str) -> Set[str]:
    """加载 tmp_question2sql.json 并提取所有 questions，返回 cleaned 模板集合"""
    seen_templates = set()
    table_stats = {}

    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
                table_key = ','.join(data.get("table", []))
                questions = data.get("questions", [])

                # 统计每个表的问题数量
                table_stats[table_key] = len(questions)

                # 提取所有 question 的模板
                for q in questions:
                    template = clean_question(q)
                    seen_templates.add(template)

            except Exception as e:
                print(f"Error parsing line {line_num}: {e}")
                continue

    # 打印统计信息
    print("=== Statistics for ./tmp_question2sql.json ===")
    total_questions = 0
    for table, count in table_stats.items():
        print(f"Table [{table}]: {count} questions")
        total_questions += count
    print(f"Total questions: {total_questions}")
    print()

    return seen_templates


def load_and_group_questions(questions_file: str) -> Dict[str, List[dict]]:
    """加载 --questions 文件，按模板分组"""
    grouped = {}

    with open(questions_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
                q_text = item["question"]
                template = clean_question(q_text)

                if template not in grouped:
                    grouped[template] = []
                grouped[template].append(item)

            except Exception as e:
                print(f"Error parsing line {line_num} in {questions_file}: {e}")
                continue

    return grouped


def find_new_question_templates(seen_templates: Set[str], grouped_questions: Dict[str, List[dict]]) -> Dict[str, List[dict]]:
    """找出未在 tmp_question2sql 中出现的模板及其问题"""
    new_groups = {}
    for template, items in grouped_questions.items():
        if template not in seen_templates:
            new_groups[template] = items
    return new_groups


def write_new_questions(new_groups: Dict[str, List[dict]], output_file: str = "tmp_newquestions.txt"):
    """将新的模板和问题写入文件"""
    with open(output_file, 'w', encoding='utf-8') as f:
        for template, items in new_groups.items():
            f.write(f"{template}\n")
            for item in items:
                f.write(f"    {json.dumps(item, ensure_ascii=False)}\n")
    print(f"New questions written to {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Evaluate and extract new SQL questions by template.")
    parser.add_argument('--questions', type=str, required=True, help='Path to the questions JSONL file')
    args = parser.parse_args()

    # Step 1: 加载 tmp_question2sql.json，获取已存在的 question templates
    seen_templates = load_tmp_questions('./tmp_question2sql.json')

    # Step 2: 加载并按模板分组新问题
    grouped_questions = load_and_group_questions(args.questions)

    # Step 3: 找出不在已知集合中的模板
    new_groups = find_new_question_templates(seen_templates, grouped_questions)

    # Step 4: 写入结果文件
    write_new_questions(new_groups)


if __name__ == '__main__':
    main()
